import os
import sys


def main():
    print("=== 大文本文件分割工具 ===")

    # 1. 用户输入要分割的文件名
    input_file = input("请输入要分隔的文本文件名（例如：data.txt）：").strip()
    if not os.path.exists(input_file):
        print(f"输入文件 {input_file} 不存在！")
        return

    # 2. 用户输入每小文件的记录数
    while True:
        try:
            batch_size = int(input("请输入每个小文件的记录数目（例如：500000）：").strip())
            if batch_size <= 0:
                raise ValueError
            break
        except ValueError:
            print("输入数据无效，请输入一个正整数！")

    # 3. 构造输出文件名模板
    base_name = os.path.splitext(input_file)[0]
    output_template = f"{base_name}_part{{}}.txt"

    print(f"\n开始分割文件 {input_file}，每 {batch_size} 行生成一个文件...")

    # 4. 逐行读取并写入小文件
    with open(input_file, 'r', encoding='utf-8', errors='ignore') as f_in:
        part_num = 1
        lines = []
        for idx, line in enumerate(f_in, 1):
            lines.append(line)
            if idx % batch_size == 0:
                output_file = output_template.format(part_num)
                with open(output_file, 'w', encoding='utf-8') as f_out:
                    f_out.writelines(lines)
                print(f"  → 已生成 {output_file}，包含 {len(lines)} 行")
                lines = []
                part_num += 1

        # 写入剩余行
        if lines:
            output_file = output_template.format(part_num)
            with open(output_file, 'w', encoding='utf-8') as f_out:
                f_out.writelines(lines)
            print(f"  → 已生成 {output_file}，包含 {len(lines)} 行")

    print("分割完成！")


if __name__ == "__main__":
    main()